# 模块化设计-反向
# 目的：训练网络，优化网络参数
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import generate
import forward

STEPS = 40000
BATCH_SIZE = 30
LEARNING_RATE_BASE = 0.001
LEARNING_RATE_DECAY = 0.999
REGULARIZER = 0.01


def backward():
    x = tf.placeholder(tf.float32, shape = (None, 2))
    y_ = tf.placeholder(tf.float32, shape = (None, 1))

    X, Y_, Y_c = generate.generateds()

    y = forward.forward(x, REGULARIZER)

    global_step = tf.Variable(0,trainable = False)

    loss_mse = tf.reduce_mean(tf.square(y-y_))
    loss_total = loss_mse + tf.add_n(tf.get_collection('losses'))
    '''
    # lossy与y_的差距，以MSE为例
    loss_mse = tf.reduce_mean(tf.square(y-y_))
    # 也可以是交叉熵和softmax的协同使用
    ce = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y,labels=tf.argmax(y_,1))
    cem = tf.reduce_mean(ce)
    # 加入正则化损失——提高泛化性
    loss_total = loss_mse + tf.add_n(tf.get_collection('losses'))
    '''

    # 指数衰减学习率——加快优化的效率
    learning_rate = tf.train.exponential_decay(
        LEARNING_RATE_BASE, # 学习率基数，初始值
        global_step, # 几轮，计数器
        300/BATCH_SIZE, # LEARNING_RATE_STEP, # 多少轮更新一次
        LEARNING_RATE_DECAY, # 衰减率
        staircase=True # False 学习率为平滑曲线
    )
    
    # 定义训练过程
    train_step = tf.train.AdamOptimizer(learning_rate).minimize(
        loss_total)


    '''
    # 滑动平均
    # 实例化滑动平均
    ema = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY,global_step)
    # 更新列表 trainable_variables自动将所有带训练数据汇总为列表
    ema_op = ema.apply(tf.trainable_variables())
    with tf.control_dependencies([train_step,ema_op]):
        train_op = tf.no_op(name='train)
    '''

    with tf.Session() as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)

        for i in range(STEPS):
            start = (i*BATCH_SIZE) % 300
            end = start + BATCH_SIZE
            sess.run(train_step, feed_dict = {x:X[start:end], y_:Y_[start:end]})
            
            # 每几轮打印出信息
            if i% 2000 == 0:
                loss_mse_v = sess.run(loss_mse, feed_dict ={x:X ,y_:Y_})
                print ("在 %d 步之后，损失达到 %f" % (i, loss_mse_v))
        # 生成 xx yy 坐标网格
        xx,yy = np.mgrid[-3:3:.01, -3:3:.01]
        # 将xx，yy拉直，并合并成一个2列的矩阵，得到网格坐标点集合
        grid = np.c_[xx.ravel(),yy.ravel()]
        probs = sess.run(y, feed_dict = {x:grid})
        probs = probs.reshape(xx.shape)

    plt.scatter(X[:,0], X[:,1], c=np.squeeze(Y_c))
    plt.contour(xx,yy,probs,levels=[.5])  # x坐标轴，y坐标轴，该点高度，等高线高度
    plt.show()

if __name__=='__main__': # 判断python运行的文件是否是主文件
    backward()

    
